package com.ruoyi.web.controller.ai;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.ruoyi.common.annotation.Anonymous;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
//import org.springframework.ai.ollama.OllamaChatModel;
//import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import com.ruoyi.common.annotation.Log;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.enums.BusinessType;
import com.ruoyi.common.core.page.TableDataInfo;
import reactor.core.publisher.Flux;

/**
 * 【会话消息】Controller
 * 
 * @author ruoyi
 * @date 2024-12-21
 */
@RestController
@RequestMapping("/ai/message")
@AllArgsConstructor
public class AiMessageController extends BaseController
{

    //AI阿里通义千问chat对象
    private final ChatModel chatModel;

    //本地大模ollama
//    private final OllamaChatModel ollamaChatModel;

    //json数据处理接口
    private final ObjectMapper objectMapper;

    /**
     * 非流式问答
     *
     * @param prompt 用户提问
     * @return org.springframework.ai.chat.model.ChatResponse
     */
    @Anonymous
    @GetMapping("chat")
    public AjaxResult chat(@RequestParam String prompt) {
        ChatClient chatClient = ChatClient.create(chatModel);
        return success(chatClient.prompt()
                // 输入单条提示词
                .user(prompt)
                // call代表非流式问答，返回的结果可以是ChatResponse，也可以是Entity（转成java类型），也可以是字符串直接提取回答结果。
                .call()
                .content());
    }

    /**
     * 流式问答
     *
     * @param prompt 用户提问
     * @return SSE流式响应
     */
    @Anonymous
    @GetMapping(value = "chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatStream(@RequestParam String prompt) {
        return ChatClient.create(chatModel).prompt()
                // 输入多条消息，可以将历史消息记录传入
                .user(prompt)
                // 流式返回
                .stream()
                // 构造SSE（ServerSendEvent）格式返回结果
                .chatResponse().map(chatResponse -> ServerSentEvent.builder(toJson(chatResponse))
                        .event("message")
                        .build()).map(ServerSentEvent::data);
    }

    /**
     * 将流式回答结果转json字符串
     *
     * @param chatResponse 流式回答结果
     * @return String json字符串
     */
    @SneakyThrows
    public String toJson(ChatResponse chatResponse) {
        return objectMapper.writeValueAsString(chatResponse);
    }

}
